package net.mycrub.bitwise.ocr.training;

import java.io.File;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

import net.mycrub.bitwise.core.Brain;
import net.mycrub.bitwise.core.util.DeserializationException;
import net.mycrub.bitwise.core.util.Random;
import net.mycrub.bitwise.core.util.SerializationException;
import net.mycrub.bitwise.ocr.reading.Sampler;

import org.apache.log4j.Logger;

public class DigitTrainer {
    private final static Logger LOGGER = Logger.getLogger(DigitTrainer.class);
    
    private final static int NB_COMBINATIONS = 100;
    private final static int NB_ITERATIONS = 1000; //per combination
    
    private final static int NB_THREADS = 3;
    
    private ListOfBrainsDAO dao;
    private List<int[]> perf = new ArrayList<int[]>();
    private Sampler sampler = new Sampler();
    
    public DigitTrainer(ListOfBrainsDAO dao) {
        this.dao = dao;
    }
    
    private synchronized void train() {
        ExecutorService pool = Executors.newFixedThreadPool(NB_THREADS);
        
        List<DigitTrainerWorker> workers = new ArrayList<DigitTrainerWorker>();
        for (int w = 0; w < NB_THREADS; w++) {
            workers.add(new DigitTrainerWorker(NB_ITERATIONS, sampler));
        }
        
        String[] allIds = dao.getAllIds();
        if (allIds.length == 0) {
            try {
                BrainFactory.createListsOfBrains(dao);
            } catch (SerializationException e) {
                // TODO Auto-generated catch block
                e.printStackTrace();
            }
            allIds = dao.getAllIds();
        }
        
        for (long i = 0; i < NB_COMBINATIONS; i++) {
            if (i > 0) {
                List<ListOfBrains> lists = new ArrayList<ListOfBrains>();
                for (DigitTrainerWorker w : workers) {
                    lists.add(w.getBrains());
                }
                combineTogether(lists, perf);
                
                for (ListOfBrains listOfBrains : lists) {
                    try {
                        dao.saveOrUpdate(listOfBrains);
                    } catch (SerializationException e) {
                        // TODO Auto-generated catch block
                        e.printStackTrace();
                    }    
                }
            }
            
            Set<String> ids = new HashSet<String>();
            for(;ids.size()<NB_THREADS;) {
                String id = allIds[(int) (Random.next() * allIds.length)];
                ids.add(id);
            }
            
            int w = 0;
            for (String id : ids) {
                try {
                    workers.get(w).setBrains(dao.get(id));
                } catch (DeserializationException e) {
                    // TODO Auto-generated catch block
                    e.printStackTrace();
                }
                w++;
            }

            List<Future<int[]>> perfs = null;
            try {
                perfs = pool.invokeAll(workers);
            } catch (InterruptedException e1) {
                // TODO Auto-generated catch block
                e1.printStackTrace();
            }
            
            int globalPerf = 0;
            perf.clear();
            for (Future<int[]> future : perfs) {
                int[] wp = null;
                try {
                    wp = future.get();
                } catch (InterruptedException e) {
                    // TODO Auto-generated catch block
                    e.printStackTrace();
                } catch (ExecutionException e) {
                    // TODO Auto-generated catch block
                    e.printStackTrace();
                }
                perf.add(wp);
                for (int p : wp) {
                    globalPerf += p;
                }
            }

            if (LOGGER.isInfoEnabled()) {
                LOGGER.info(i + "/" + NB_COMBINATIONS + " : " + globalPerf);
            }
        }
    }

    private void combineTogether(List<ListOfBrains> lists, List<int[]> perfs) {
        List<Brain> brains = new ArrayList<Brain>();
        int[] perf = new int[perfs.size() * perfs.get(0).length]; //XXX ugly
        for(int i = 0; i < lists.size(); i++) {
            brains.addAll(lists.get(i).list);
            System.arraycopy(perfs.get(i), 0, perf, perfs.get(i).length * i, perfs.get(i).length);
        }
        
        Set<CoupleCandidate> ccs = new HashSet<CoupleCandidate>();
        int size = brains.size();
        
        int nbCandidates = (int) Math.pow(brains.size(), 1.5);
            
        while(ccs.size() < nbCandidates) {
            int mom = (int)(brains.size() * Random.next());
            int dad = (int)(brains.size() * Random.next());
            ccs.add(new CoupleCandidate(
                    brains.get(mom), 
                    brains.get(dad), 
                    perf[mom]*perf[dad]));
        }
        
        List<CoupleCandidate> sorted = new LinkedList<CoupleCandidate>(ccs);
        Collections.sort(sorted);
        
        brains.clear();
        Iterator<CoupleCandidate> it = sorted.iterator();
        for (int i = 0; i < lists.size() * lists.get(0).list.size(); i++) {
            CoupleCandidate next = it.next();
            brains.add(next.getMom().combine(next.getDad()));
        }
        
        for(ListOfBrains lob : lists) {
            int bsize = lob.list.size();
            lob.list.clear();
            for(int i = 0; i < bsize; i++) {
                lob.list.add(brains.remove(0));
            }
        }
    }
    
    private static void usage() {
        System.err.println("Usage: DigitTrainer brains-directory");
        System.err.println("       brains-directory must be a directory");
        System.exit(1);
    }
    
    public static void main(String[] args) {
        if (args.length < 1) {
            usage();
        }

        String dirname = args[0];
        File dir = new File(dirname);
        
        if(!dir.exists() || !dir.isDirectory()) {
            usage();
        }

        ListOfBrainsDAO dao = new ListOfBrainsDAO();
        dao.setDir(dir);
        
        new DigitTrainer(dao).train();
    }
}
